"""
该文件仅用于定义跨文件使用的数据集处理方法
"""
import os
import random
import numpy as np
import config


SEQ_MAX_LEN = config.SEQ_MAX_LEN
SEPARATE_RATE = config.SEPARATE_RATE
FILE_DATA = config.FILE_DATA
FILE_DATA_SHUFFLE = config.FILE_DATA_SHUFFLE
FILE_DATA_TRAIN = config.FILE_DATA_TRAIN
FILE_DATA_TEST = config.FILE_DATA_TEST


# 对txt文件按行shuffle
def shuffle_file_txt(before_shuffle, after_shuffle):
    before_shuffle = open(before_shuffle, 'rU')
    list_shuffle = []
    for line in before_shuffle:
        list_shuffle.append(line)
    random.shuffle(list_shuffle)
    shuffle_file = after_shuffle
    with open(shuffle_file, 'w') as after_shuffle:
        for index in range(len(list_shuffle)):
            after_shuffle.write(list_shuffle[index])
    before_shuffle.close()
    after_shuffle.close()


# 分割数据集
def separate_data(data_shuffle, data_train, data_test, percentage):
    data_shuffle = open(data_shuffle, 'rU')
    list_data = []
    for line in data_shuffle:
        list_data.append(line)
    line_separate = int(percentage * len(list_data))
    with open(data_train, 'w') as data_train:
        for index in range(line_separate):
            data_train.write(list_data[index])
    with open(data_test, 'w') as data_test:
        for index in range(line_separate, len(list_data)):
            data_test.write(list_data[index])
    data_shuffle.close()
    data_train.close()
    data_test.close()


# 读取txt中保存的数据，并返回一个dic
def read_data_txt(file_data):
    file_read = open(file_data, 'rU')
    data_number = len(file_read.readlines())
    file_read.seek(0)
    data_x = np.zeros([data_number, SEQ_MAX_LEN, 5], dtype='float32')   # 实际数据，其中包含五个维度
    data_y = np.zeros([data_number, 1], dtype='int32')   # 数据标签
    data_len = []   # 数据长度
    line_count = 0  # 行计数器
    for line in file_read:
        line = line.strip('\n')     # 去除行末换行符
        set_data = line.split(';')  # 按文件内定义格式分割
        set_ang0 = set_data[0]  # 指尖0的角度数据
        set_ang1 = set_data[1]  # 指尖1的角度数据
        set_ang2 = set_data[2]  # 指尖2的角度数据
        set_ang3 = set_data[3]  # 指尖3的角度数据
        set_ang4 = set_data[4]  # 指尖4的角度数据
        set_label = set_data[5]     # 数据标签
        # 获取单个数据
        single_ang0 = set_ang0.split()  # 单个指尖0角度数据
        single_ang1 = set_ang1.split()  # 单个指尖1角度数据
        single_ang2 = set_ang2.split()  # 单个指尖2角度数据
        single_ang3 = set_ang3.split()  # 单个指尖3角度数据
        single_ang4 = set_ang4.split()  # 单个指尖4角度数据
        data_len.append(len(single_ang0))
        data_y[line_count][0] = float(set_label)
        for index in range(int(len(single_ang0))):
            data_x[line_count][index][0] = single_ang0[index]
            data_x[line_count][index][1] = single_ang1[index]
            data_x[line_count][index][2] = single_ang2[index]
            data_x[line_count][index][3] = single_ang3[index]
            data_x[line_count][index][4] = single_ang4[index]
        line_count += 1
    file_read.close()
    dic_data = {'x': data_x, 'y': data_y, 'length': data_len, 'number': data_number}
    return dic_data     # 返回一个包含多个参数的dictionary


# 读取list中的数据，并返回一个dic
def read_data_list(angle_list, distance_list, slice_list, data_number):
    data_x = np.zeros([data_number, SEQ_MAX_LEN, 2], dtype='float32')   # 实际数据，其中包含两个维度
    data_sta = np.zeros([data_number, 1], dtype='int32')  # 数据开始序号
    data_end = np.zeros([data_number, 1], dtype='int32')    # 数据结束序号
    data_len = []   # 数据长度
    for i in range(len(slice_list)):
        line_ang = angle_list[slice_list[i][0]: slice_list[i][1]]
        line_dis = distance_list[slice_list[i][0]: slice_list[i][1]]
        for j in range(len(line_ang)):
            data_x[i][j][0] = line_ang[j]
            data_x[i][j][1] = line_dis[j]
        data_sta[i] = slice_list[i][0]
        data_end[i] = slice_list[i][1]
        data_len.append(len(line_ang))
    dic_data = {'x': data_x, 'length': data_len,
                'number': data_number, 'start': data_sta, 'end': data_end}
    return dic_data


if __name__ == '__main__':
    # 删除上次训练产生的临时文件
    if os.path.exists(FILE_DATA_SHUFFLE):
        os.remove(FILE_DATA_SHUFFLE)
    if os.path.exists(FILE_DATA_TRAIN):
        os.remove(FILE_DATA_TRAIN)
    if os.path.exists(FILE_DATA_TEST):
        os.remove(FILE_DATA_TEST)
    shuffle_file_txt(FILE_DATA, FILE_DATA_SHUFFLE)    # 数据集shuffle
    separate_data(FILE_DATA_SHUFFLE, FILE_DATA_TRAIN, FILE_DATA_TEST, SEPARATE_RATE)   # 数据集分割
    print('Data processing completed!')
